#include<unistd.h>

#include<chrono>
#include<cstring>
#include<fstream>
#include<iostream>
#include<map>
#include<vector>
#include "gmem_ascend.h"
#include "acl/acl.h"

using namespace std;
using namespace chrono;

#define SUCCESS 0
#define FAILED 1
#define TWO_MB (2*1024*1024)

#define INFO_LOG(fmt, args...) fprintf(stdout, "[INFO] " fmt "\n", ##args)
#define WARN_LOG(fmt, args...) fprintf(stdout, "\033[31m[WARN]" fmt "\033[0m\n", ##args)
#define ERROR_LOG(fmt, args...) fprintf(stdout, "\033[1;31m[ERROR]" fmt "\033[0m\n", ##args)

int64_t N = 4096;
int gmem_fd = 0;
int device_id = 0;
int sleep_time = 1;
aclError ret, ret1;
pthread_t thread;
bool InitResource()
{
    if(aclInit(nullptr) != ACL_SUCCESS){
        cout << "Init acl failed\n";
        return false;
    }
    if(aclrtSetDevice(device_id) != ACL_SUCCESS){
        cout << "Set device failed. deviceId is " << device_id << endl;
        return false;
    }
    if(aclopSetModeDir("./model") != ACL_SUCCESS){
        cout << "Load single op model failed" << endl;
        return false;
    }
    return  true;
}
void DestoryResource()
{
    bool flag=false;
    if(aclrtResetDevice(device_id) != ACL_SUCCESS){
        ERROR_LOG("Reset device %d failed",device_id);
        flag = true;
    }
    if(aclFinalize() != ACL_SUCCESS){
        ERROR_LOG("Finalize acl failed");
        flag = true;
    }
    if(flag){
        ERROR_LOG("Destory resource failed");
    } else {
        INFO_LOG("Destory resource success");
    }
}

//申请新的 2MB 区域
void *allocate_new_memory(void *arg){
    sleep(10);
    printf("Thread 1:Allocating new 2MB memory...\n");
    void *new_memory = NULL;
    for (int i = 0; i < 5120; i++){
        new_memory = mmap(NULL, TWO_MB, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_PERR_SHARED, -1, 0);
        if (new_memory == MAP_FAILED){
            perror("Thread 2: mmap failed");
            exit(1);
        }
        // printf("Thread 2: New memory %d allocated at %p\n", i, new_memory);
        if(munmap(new_memory, TWO_MB) == -1){
            perror("munmap failed for new_memory");
            break
        }
    }
    printf("Thread 1:Allocating done...\n");
    return NULL;
}

int NormalRun()
{
    printf("Matmul: running...\n");
    // 注册stream
    aclrtStream stream1 = nullptr;
    ret = aclrtCreateStream(&stream1);

    void *device_A, *device_B, *device_C;
    size_t matrix_size = N * N * sizeof(aclFloat16);

    device_A = (aclFloat16 *)mmap(NULL, matrix_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_PERR_SHARED, -1, 0);
    device_B = (aclFloat16 *)mmap(NULL, matrix_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_PERR_SHARED, -1, 0);
    device_C = (aclFloat16 *)mmap(NULL, matrix_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_PERR_SHARED, -1, 0);

    int64_t dim[2] = {N, N};
    aclTensorDesc *input_desc[2], *output_desc[1];
    input_desc[0] = aclCreateTensorDesc(ACL_FLOAT16, 2, dim, ACL_FORMAT_ND);
    input_desc[1] = aclCreateTensorDesc(ACL_FLOAT16, 2, dim, ACL_FORMAT_ND);
    output_desc[0] = aclCreateTensorDesc(ACL_FLOAT16, 2, dim, ACL_FORMAT_ND);

    aclDataBuffer *input_data[2], *output_data[1];
    inputs[0] = aclCreateDataBuffer(device_A, matrix_size);
    inputs[1] = aclCreateDataBuffer(device_B, matrix_size);
    outputs[0] = aclCreateDataBuffer(device_C, matrix_size);

    aclopAttr *attr = aclCreateDataBuffer();
    aclopSetAttrBool(attr, "transpose_x1", 0);
    aclopSetAttrBool(attr, "transpose_x2", 0);

    fill((aclFloat16 *)device_A, (aclFloat16 *)device_A + matrix_size / sizeof(aclFloat16), aclFloatTofloat16(1));//写入数据
    fill((aclFloat16 *)device_B, (aclFloat16 *)device_B + matrix_size / sizeof(aclFloat16), aclFloatTofloat16(1));

    auto start = system_clok::now();
    printf("Matmul: execute...\n");
    ret = aclopExecuteV2("MatMul", 2, input_desc, inputs, 1, output_desc, outputs, attr, stream1);//加载执行算子 dev侧会触发缺页
    if(ret != ACL_SUCCESS){
        ERROR_LOG("Launch MatMul failed. errorCode is %d", static_cast<int32_t>(ret));
    }
    ret = aclrtSynhronizeDevice(); // 同步，等待device运行完成
    if (ret != ACL_SUCCESS){
        ERROR_LOG("Execute MatMul failed. errorcode is %d",static_cast<int32_t>(ret));
    }
    auto end = system_clock::now();
    // INFO_LOG("Launch finish, Press Enter to Continue.");
    // getchar();
    auto timeUs = duration_cast<microseconds>(end - start).count();

    INFO_LOG("check result %f", aclFloat16ToFloat(*(aclFloat16 *)device_C));
    cout << "time: "<<timeUs << "us" << endl;
    munmap(device_A, matrix_size);
    munmap(device_B, matrix_size);
    munmap(device_C, matrix_size);
    aclDestroyDataBuffer(inputs[0]);
    aclDestroyDataBuffer(inputs[1]);
    aclDestroyDataBuffer(outputs[0]);
    aclrtDestroyStream(stream1);
    printf("Matmul: done...\n");
    return NULL;
}
#define SIZE_1G (1024 * 1024 * 1024)
#define SIZE_25G ((unsigned long) 25 * SIZE_1G)

int mian(int argc, char *argv[])
{
    int ret;
    int hnid = 0;
    if(argc >= 2){
        sscanf(argv[1], "%ld", &N);
    }
    void *fill_device;
    // acl初始化
    if (!InitResource()){
        ERROR_LOG("Init resource failed");
        return FAILED;
    }
    INFO_LOG("Init resource success");

    hnid = get_hnid(devivce_id);
    if(hnid < 0){
        return 1;
    }

    INFO_LOG("prefill device memory, Press Enter to Continue.");
    getchar();

    //创建线程来模拟并发
    pthread_t fault_thread, alloc_thread;

    if(pthread_create(&alloc_thread, NULL, allocate_new_memory, NULL) ! = 0){
        perror("Failed to create thread 2");
        ret = EXIT_FAILURE;
        goto out;
    }

    NormalRun();

    // //等待线程结束
    // pthread_join(fault_thread, NULL);
    pthread_join(alloc_thread, NULL);

out:
    // acl去初始化
    // munmap(fill_device, (unsigned long)25 * SIZE_1G);
    DestoryResource();
    return 0;
}

